import pandas as pd
import pymysql
from tqdm import tqdm


def train_data(code, type="index"):
    yk = pymysql.connect(
        host='10.168.10.50',
        port=3306,
        user='root',
        password='123456',
        database='youkuang',
        charset='utf8'
    )

    sentiment = pymysql.connect(
        host='10.168.10.34',
        port=3306,
        user='root',
        password='root',
        database='zl_sentiment',
        charset='utf8'
    )

    hl = pd.read_sql(
        "select openPrice as openHl,closePrice as closeHl,highPrice as highHl,lowPrice as lowHl,date as tradeDate from changerate where changeType='cnyusd'",
        con=sentiment)
    hl["tradeDate"] = hl["tradeDate"].dt.strftime('%Y-%m-%d')

    lv = pd.read_sql(
        "select zm3,zm6,z1,z2,z3,z5,z7,z10,z30,m2,m5,m10,m30,time as tradeDate from bondyield",
        con=sentiment)
    lv["tradeDate"] = pd.to_datetime(lv['tradeDate'])
    lv["tradeDate"] = lv["tradeDate"].dt.strftime('%Y-%m-%d')

    value_data = None
    if type == "index":
        value_data = pd.read_sql(
            "SELECT open_index AS openIndex,close_index AS closeIndex,highest_index AS highIndex,lowest_index AS lowIndex,turnover_value as turnoverValueIndex,turnover_vol as turnoverVolIndex,trade_date AS tradeDate FROM mkt_idxd WHERE ticker_symbol='" + code + "'",
            con=yk)
        value_data["tradeDate"] = pd.to_datetime(value_data['tradeDate'])
        value_data["tradeDate"] = value_data["tradeDate"].dt.strftime('%Y-%m-%d')
    elif type == "stock":
        value_data_piece_1 = pd.read_sql(
            "SELECT mkt_equd_adj_af.OPEN_PRICE_2 AS openIndex, mkt_equd_adj_af.CLOSE_PRICE_2 AS closeIndex, mkt_equd_adj_af.HIGHEST_PRICE_2 AS highIndex, mkt_equd_adj_af.LOWEST_PRICE_2 AS lowIndex, mkt_equd_adj_af.trade_date AS tradeDate FROM mkt_equd_adj_af USE INDEX (TICKER_SYMBOL) WHERE mkt_equd_adj_af.TICKER_SYMBOL='" + code + "'",
            con=yk)
        value_data_piece_1["tradeDate"] = pd.to_datetime(value_data_piece_1['tradeDate'])
        value_data_piece_1["tradeDate"] = value_data_piece_1["tradeDate"].dt.strftime('%Y-%m-%d')
        value_data_piece_2 = pd.read_sql(
            "SELECT mkt_equd.turnover_value AS turnoverValueIndex, mkt_equd.turnover_vol AS turnoverVolIndex, mkt_equd.trade_date AS tradeDate FROM mkt_equd USE INDEX (TICKER_SYMBOL) WHERE mkt_equd.TICKER_SYMBOL='" + code + "' and turnover_vol>0",
            con=yk)
        value_data_piece_2["tradeDate"] = pd.to_datetime(value_data_piece_2['tradeDate'])
        value_data_piece_2["tradeDate"] = value_data_piece_2["tradeDate"].dt.strftime('%Y-%m-%d')
        value_data = pd.merge(value_data_piece_1, value_data_piece_2, on="tradeDate")

    temp = pd.merge(value_data, lv, on="tradeDate")
    result = pd.merge(temp, hl, on="tradeDate")
    result["profit"] = result['closeIndex'] > result['closeIndex'].shift(1)
    result["profitPct"] = result['closeIndex'] / result['closeIndex'].shift(1) - 1
    result = result.sort_values(by="tradeDate")
    result.to_csv(code + type + ".csv", index=False)


def stock_list():
    yk = pymysql.connect(
        host='10.168.10.50',
        port=3306,
        user='root',
        password='123456',
        database='youkuang',
        charset='utf8'
    )
    stock = pd.read_sql(
        "SELECT mkt_equd_adj_af.TICKER_SYMBOL AS ticker FROM mkt_equd_adj_af WHERE mkt_equd_adj_af.TRADE_DATE='2023-11-22'",
        con=yk)
    return [x for x in stock["ticker"].values.tolist() if x.startswith(('3', '6', '0'))]


data = stock_list()
for stock in tqdm(data, desc="股票列表"):
    train_data(stock, type="stock")
